import os
import matplotlib.pyplot as plt
import mindspore.dataset as ds
import mindspore.dataset.vision.c_transforms as cvision


def MNIST(root, usage="train", sampler=None, transform=None, download=False):
    """ MNIST Dataset.
    Args:
        root (string): Root directory of dataset where ``*-ubyte.gz`` exist.
        usage (string, optional): If True, creates dataset from training part,
            otherwise from test part.
        sampler (SamplerObj, optional): Object used to choose samples from dataset.
        transform (callable, optional): A function/transform that takes in an image
            and returns a transformed version. E.g, ``dataset.vision.RandomCrop``
        download (bool, optional): If true, downloads the dataset from the internet and
            puts it in root directory. If dataset is already downloaded, it is not
            downloaded again.
    """
    resources1 = ["http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz",
                  "http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz",
                  "http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz",
                  "http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz"]

    resources2 = ["t10k-images-idx3-ubyte",
                  "t10k-labels-idx1-ubyte",
                  "train-images-idx3-ubyte",
                  "train-labels-idx1-ubyte"]

    def check_exist(rt_path, resources):
        for res in resources:
            if not os.path.exists(os.path.join(rt_path, res)):
                return False
        return True

    def download():
        if check_exist(root, [res.split("/")[-1] for res in resources1]):
            return
        if check_exist(root, resources2):
            return
        os.makedirs(root, exist_ok=True)
        for res in resources1:
            if not os.path.exists(os.path.join(root, res.split("/")[-1])):
                os.system("wget -P " + root + " " + res)

    if download:
        download()

    if check_exist(root, [res.split("/")[-1] for res in resources1]):
        for res in resources1:
            os.system("gzip -d " + root + "/" + res.split("/")[-1])
    elif check_exist(root, resources2):
        pass
    else:
        raise RuntimeError('Dataset not found. You can use download=True to download it')

    # mindspore built-in api
    dataset = ds.MnistDataset(dataset_dir=root,          # dir hold mnist dataset
                              usage=usage,               # "train"/"val"/"all" part of dataset
                              sampler=sampler,           # sampler
                              num_parallel_workers=4)    # multithreading

    if transform:
        if not isinstance(transform, list):
            transform = [transform]
        dataset = dataset.map(operations=transform, input_columns="image")

    return dataset


if __name__ == '__main__':
    dataset_dir = "./dataset/mnist"
    mnist = MNIST(dataset_dir, download=True, transform=[cvision.Resize(40), cvision.RandomCrop(28)])
    
    # set random seed for shuffle
    #ds.config.set_seed(1)

    # To see what samples do we get from dataset
    for index, data in enumerate(mnist.create_dict_iterator(output_numpy=True)):
        if index >= 10:
            break
        print(data["image"].shape, data["label"])
        plt.subplot(2, 5, index+1)
        plt.imshow(data["image"].squeeze(), cmap=plt.cm.gray)
        plt.title(data["label"])
    plt.show()
